import pandas as pd
import torch
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
from PIL import Image
import base64
import numpy as np
import logging
import os
import datetime
from pathlib import Path



data_path = r'/mnt/6t/lyh/1/M2E2/M2E2/text/mee.tsv'

data = pd.read_csv(data_path, sep='\t')
# model_path = r'/home/hpuser/home/hzy/model/QWen2-7B'
model_path = '/mnt/6t/lyh/1/Qwen2-VL-7B-Instruct'
# model_path = r'/home/student/hzy/hzy/model/Qwen2-VL-7B-Instruct'
# model_path = '/home/student/hzy/hzy/model/Qwen-Lora-MNRE/test1'
model = Qwen2VLForConditionalGeneration.from_pretrained(model_path, torch_dtype="auto", device_map="auto")
processor = AutoProcessor.from_pretrained(model_path)


# Step 3: 读取图像
# def load_image(image_path):
#     image = Image.open(image_path).convert("RGB")
#     return image

def convert_to_nan(value):
    if isinstance(value, str) and value.lower() == 'nan':
        return np.nan
    return value

def local_ner_predict(text, image_path):
    with open(image_path, 'rb') as img_file:
        img_data = img_file.read()
    img_data_b64 = base64.b64encode(img_data).decode('utf-8')


    messages = [
        # {
        #     "role": "system",
        #     "content": "你是一名多模态关系抽取领域的专家。任务是从给定文本、图像和实体对（头实体 h 和尾实体 t）中判断两者关系。1. 仅针对头实体 h 和尾实体 t：不要更改或新增其他实体。2.如果图像和文本没有因果关联，关系输出为error，输出格式为('h', 'error', 't')，如('Crocodile', 'error', 'South Wales')。2. 关系选项：held_on、couple、alternate_names、member_of、peer、contain、nationality、part_of、locate_at、subsidiary、present_in、charges、parent、place_of_residence、awarded、siblings、place_of_birth、religion、neighbor、nan。只允许从这些关系选项中进行选择，绝对不允许生成其它关系。3. 输出格式：结果必须是三元组('h', 'relation', 't')，如('Crocodile', 'held_on', 'South Wales')。若无关系，输出('h', 'nan', 't')，不要输出('nan','nan','nan')。充分利用图像和文本信息，严格按要求抽取关系。"
        # },
        # {
        #     "role": "system",
        #     "content":"你是一名多模态关系抽取领域的专家。任务是从给定文本、图像和实体对（头实体 h 和尾实体 t）中判断两者关系。首先判断文本和图像是否一一对应.如果文本和图像没有因果关联,输出格式为('h', 'error', 't')，如('Crocodile', 'error', 'South Wales')。如果文本和图像有因果关联，但是实体之间没有关系，输出格式为('h', 'nan', 't')，('Crocodile', 'nan', 'South Wales')。如。2. 仅针对头实体 h 和尾实体 t：不要更改或新增其他实体。3. 关系选项：held_on、couple、alternate_names、member_of、peer、contain、nationality、part_of、locate_at、subsidiary、present_in、charges、parent、place_of_residence、awarded、siblings、place_of_birth、religion、neighbor。只允许从这些关系选项中进行选择，绝对不允许生成其它关系。4. 输出格式：结果必须是三元组('h', 'relation', 't')，如('Crocodile', 'held_on', 'South Wales')。不要输出('nan','nan','nan')。充分利用图像和文本信息，严格按要求抽取关系。"
        # },
        # {
        #     "role": "system",
        #     "content":"你是一名多模态关系抽取领域的专家。首先判断文本和图像是否一一对应，再详细理解文本个图像的内容，之后根据两者共有的关系从给定文本、图像和实体对（头实体 h 和尾实体 t）中判断两者关系。1. 仅针对头实体 h 和尾实体 t：不要更改或新增其他实体。2. 关系选项：contain、place_of_birth、none、charges、peer、locate_at、member_of、alternate_names、present_in、colleague、part_of、creat、subsidiary、capital、place_of_residence、place_of_death、couple、awarded、administration_divisions、leader、held_on、nationality、parent、founder_of、siblings、relative、neighbor。严格从这些关系选项中进行选择，绝对不允许生成其它关系。3. 输出格式：结果必须是三元组('h', 'relation', 't')，如('Crocodile', 'held_on', 'South Wales')。4.若无关系，输出('h', 'none', 't')，不要输出('none','none','none')。充分利用图像和文本信息，严格按要求抽取关系。"
        # },
        # top 1
        # {
        #      "role": "system",
        #     "content":"请仔细阅读以下文本和观察图像内容。对于给定的头实体 h 和尾实体 t，请根据文本数据和图像数据辅助，判断头实体和尾实体的关系，有以下关系选项：contain（包含关系）、place_of_birth（出生地）、none（没有明确关系）、charges（指控或罪名）、peer（同级或同类关系）、locate_at（位于）、member_of（成员或隶属关系）、alternate_names（别名或变体）、present_in（出现关系）、colleague（同事关系）、part_of（整体与部分关系）、creat（创建关系）、subsidiary（子公司关系）、capital（首都关系）、place_of_residence（居住地）、place_of_death（死亡地）、couple（夫妻或伴侣关系）、awarded（获奖关系）、administration_divisions（行政区划关系）、leader（领导关系）、held_on（事件发生或举行地）、nationality（国籍关系）、parent（母公司或上级组织关系）、founder_of（创始人关系）、siblings（兄弟姐妹关系）、relative（亲属关系）、neighbor（邻近或相邻关系）、none（没有关系）。从文本和图像中抽取两者共有的关系。注意，仅针对头实体 h 和尾实体 t 进行判断，不要更改或新增其他实体。最后，以三元组 ('h', 'relation', 't') 的形式输出结果，如 ('Crocodile', 'held_on', 'South Wales')。如果无法确定关系，则输出 ('h', 'none', 't')。"

        # },
        {
            "role": "system",
            "content":"请仔细阅读以下文本和观察图像内容。请根据文本数据和图像数据辅助进行事件抽取，有以下事件抽取选项：'Justice:Arrest-Jail','Conflict:Attack','Movement:Transport','Contact:Phone-Write','Life:Die','Conflict:Demonstrate','Transaction:Transfer-Money','Contact:Meet','none'。注意，不要更改或新增其他事件选项。请在输出时仔细分析文本和图片数据。输出时直接输出选项，不需要添加其他符号"
        },
        {
            "role": "user",
            "content": [
                # {"type": "text", "text": "请先对以下数据进行详细分析，再进行事件抽取任务,输出时以给出的选项为主."},
                {"type": "text", "text": text},
                # {"type": "text", "text": text1},
                {"type": "image_path", "image_path": image_path}
            ]
        }
    ]

    chat_template = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, video_inputs = process_vision_info(messages)

    inputs = processor(
        text=[chat_template],
        images=image_inputs,
        # videos=video_inputs,
        padding=True,
        return_tensors="pt"
    ).to(model.device)
    
        # 清空缓存，避免显存过度增长
    torch.cuda.empty_cache()

    generated_ids = model.generate(**inputs, max_new_tokens=256)
    generated_ids_trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]

    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    print(output_text)
    return output_text[0]


# 初始化计数器
TP = 0
FP = 0
FN = 0
# 定义允许的关系选项列表
allowed_relations = [
'Justice:Arrest-Jail',
'Conflict:Attack',
'Movement:Transport',
'Contact:Phone-Write',
'Life:Die',
'Conflict:Demonstrate',
'Transaction:Transfer-Money',
'Contact:Meet',
'none'
]

# 创建字典来跟踪每个事件类型的TP、FP、FN
event_stats = {}
for relation in allowed_relations:
    event_stats[relation] = {"TP": 0, "FP": 0, "FN": 0, "support": 0}

for i, row in data.iterrows():
    print("\n--- Example ---")
    image_path = f"/mnt/6t/lyh/1/M2E2/M2E2/image/image/image/{row['#1 ImageID']}"
    print(image_path)
    re_result = local_ner_predict(row['#2 String'], image_path)
    ttext = row['#2 String']
    print(ttext)
    true_ee = row['#3 String']
    print("真实事件是:", true_ee)
    print("预测的事件:", re_result)
    
    # 更新真实事件的支持数
    if true_ee in event_stats:
        event_stats[true_ee]["support"] += 1

    try:
        # 将预测结果转为三元组
        # triplet = ast.literal_eval(re_result)
        # re_relation = triplet[1]  # 第二个元素是关系
        re_relation = re_result
        re_relation = convert_to_nan(re_relation)
        # print("抽取的事件是:", re_relation)  # 例如：couple

        # 1. 如果预测值与真实值相等
        if re_relation == true_ee:
            TP += 1
            if re_relation in event_stats:
                event_stats[re_relation]["TP"] += 1

        # 2. 如果预测值不为none且与真实值不相等
        elif re_relation != 'none':
            FP += 1
            if re_relation in event_stats:
                event_stats[re_relation]["FP"] += 1
            
            # 如果真实值不为none，则真实值对应的类别FN+1
            if true_ee != 'none' and true_ee in event_stats:
                event_stats[true_ee]["FN"] += 1

        # 3. 如果预测值为none且真实值不为none
        elif re_relation == 'none' and true_ee != 'none':
            FN += 1
            if true_ee in event_stats:
                event_stats[true_ee]["FN"] += 1

        # 输出当前累计的TP、FP、FN数量
        print(f"当前累计 - TP: {TP}, FP: {FP}, FN: {FN}")

    except Exception as e:
        print("解析预测结果时发生错误:", e)
        FN += 1  # 如果解析出错，计为FN
        if true_ee != 'none' and true_ee in event_stats:
            event_stats[true_ee]["FN"] += 1
        print(f"当前累计 - TP: {TP}, FP: {FP}, FN: {FN}")

# 计算精确度、召回率和F1值
precision = TP / (TP + FP) if TP + FP > 0 else 0
recall = TP / (TP + FN) if TP + FN > 0 else 0
f1_score = (2 * precision * recall) / (precision + recall) if precision + recall > 0 else 0

# 输出结果
print("\n--- 结果统计 ---")
print(f"TP (True Positives): {TP}")
print(f"FP (False Positives): {FP}")
print(f"FN (False Negatives): {FN}")
print(f"精确度 (Precision): {precision:.4f}")
print(f"召回率 (Recall): {recall:.4f}")
print(f"F1 值 (F1-Score): {f1_score:.4f}")

# 输出每个事件类型的统计信息
print("\n--------------------------------------------------------------------------------")
print("{:<25} {:<10} {:<10} {:<10} {:<10}".format("事件类型", "精确度", "召回率", "F1值", "支持数"))
print("--------------------------------------------------------------------------------")
total_support = 0
macro_precision = 0
macro_recall = 0
macro_f1 = 0
weighted_precision = 0
weighted_recall = 0
weighted_f1 = 0
valid_event_count = 0

for event, stats in event_stats.items():
    event_precision = stats["TP"] / (stats["TP"] + stats["FP"]) if (stats["TP"] + stats["FP"]) > 0 else 0
    event_recall = stats["TP"] / (stats["TP"] + stats["FN"]) if (stats["TP"] + stats["FN"]) > 0 else 0
    event_f1 = 2 * event_precision * event_recall / (event_precision + event_recall) if (event_precision + event_recall) > 0 else 0
    
    # 只计算有支持数的事件类型
    if stats["support"] > 0:
        valid_event_count += 1
        macro_precision += event_precision
        macro_recall += event_recall
        macro_f1 += event_f1
        
    total_support += stats["support"]
    weighted_precision += event_precision * stats["support"]
    weighted_recall += event_recall * stats["support"]
    weighted_f1 += event_f1 * stats["support"]
    
    print("{:<25} {:<10.4f} {:<10.4f} {:<10.4f} {:<10}".format(
        event, event_precision, event_recall, event_f1, stats["support"]))

# 计算宏平均和加权平均
if valid_event_count > 0:
    macro_precision /= valid_event_count
    macro_recall /= valid_event_count
    macro_f1 /= valid_event_count

if total_support > 0:
    weighted_precision /= total_support
    weighted_recall /= total_support
    weighted_f1 /= total_support

print("--------------------------------------------------------------------------------")
print("{:<25} {:<10.4f} {:<10.4f} {:<10.4f} {:<10}".format("宏平均", macro_precision, macro_recall, macro_f1, total_support))
print("{:<25} {:<10.4f} {:<10.4f} {:<10.4f} {:<10}".format("加权平均", weighted_precision, weighted_recall, weighted_f1, total_support))
print("================================================================================")